import os
import numpy as np
import pandas as pd
import requests
from bs4 import BeautifulSoup
"""
CG99模型网解析：
一、主网址cg99设计网https://www.cg99.cn/models/model/model-{}-0-0-0-1-0-0-0-0-0-0-{}-0-0-0-0-0.html.
类别：全部 人物 动物 场景 武器 机械 生活用品 车辆交通 文物雕塑 植物山石 科学 飞机船艇 电子家电 美陈 展厅 建筑 室内家具
0.全部
1.人物
2.动物


存储结构：
按着爬的sort即可进行分类,sort映射中文名称
按page的数组即可进行分页下载，可以进行间断性下载
已原图片名称命名图片，整个page目录下同时生成对应的映射csv文件
所以目录结构如下：
/sort 类别名称
     /page 页面号
          /图片png  page.csv该页的图片信息表
          
"""


class CGcla99(object):

    def __init__(self):
        super(CG99, self).__init__()
        self.header = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/72.0.3626.81 Safari/537.36'} # 设置http header，视情况加需要的条目，这里的token是用来鉴权的一种方式

        self.sort_map={'1':'人物','2':'动物','3':'场景','4':'武器','5':'机械','6':'生活用品','7':'车辆交通','8':'文物雕塑',
                       '9':'植物山石','10':'科学','11':'飞机船艇','12':'电子家电','13':'美陈','14':'展厅','15':'建筑',
                       '16':'室内家具'}


    def ex_text(self,url,sort_num,page_num,save_dir):
        sort_type = self.sort_map[str(sort_num)]
        res = requests.get(url)
        sp=BeautifulSoup(res.text)
        text = sp.find_all("a", "top-cover")
        # print(text)
        data_list=[]
        for i in text:
            #     print(i)
            labels=i['title']
            img_url=i.img['src'] if not i.img['src'].startswith('data') else i.img['data-url']
            # img_path=img_url.split('!')[0]
            # print(type(i.a.img['alt']))
            # print(i.a.img['data-src'])
            img_name=os.path.split(img_url)[1].split('!')[0]
            data_list.append([img_name,labels])
            self.download_img(sort_type,page_num,img_url,save_dir)
        # print(data_list)
        return data_list


    def download_img(self,sort_type,page_num,img_url,save_dir):
        """
        :param sort_type: 模型类别
        :param page_num: 页面号
        :param img_url: 传入图片url
        :param save_dir: 保存的总目录
        :return:
        """
        img_path = img_url.split('!')[0]
        img=os.path.split(img_path)[1]
        # img_name=os.path.splitext(img)[0]
        img_dir=os.path.join(save_dir,sort_type,str(page_num))
        if not os.path.exists(img_dir):
            os.makedirs(img_dir)
        save_path=os.path.join(img_dir,img)
        r = requests.get(img_url, headers=self.header, stream=True)
        if r.status_code == 200:
            with open(save_path,'wb') as w:
                w.write(r.content)


    def download(self,save_dir,sort_num,start_page,end_page):
        """
        :param save_dir:   保存的目录
        :param sort_num:  类型索引,从1开始到24
        :param all_pages:  需要爬的总页面数
        :return:
        """
        sort_type=self.sort_map[str(sort_num)]
        for i in range(start_page,end_page):
            url = 'https://www.cg99.cn/models/model/model-{}-0-0-0-1-0-0-0-0-0-0-{}-0-0-0-0-0.html'.format(sort_num, i)
            data_list=self.ex_text(url,sort_num,i,save_dir)
            img_dir = os.path.join(save_dir, sort_type, str(i))
            csv_name=str(i)+'.csv'
            csv_path=os.path.join(img_dir,csv_name)
            self.save_csv(csv_path,data_list)

    def save_csv(self,csv_path,data_list):
        """
        将该页的所有模型信息csv存下
        :param csv_path:  本地页面的绝对路径
        :param data_list:  数据集
        :return:
        """
        ar=np.array(data_list)
        df=pd.DataFrame(ar,columns=['imgName','label'],index=range(1,len(data_list)+1))
        df.to_csv(csv_path)


if __name__ == '__main__':
    cg99=CG99()
    url='https://www.cg99.cn/models/model/model-1-0-0-0-1-0-0-0-0-0-0-2-0-0-0-0-0.html'
    save_dir='imgs'
    # cg99.ex_text(url,1,2,save_dir)
    cg99.download(save_dir,1,1,8)
    # cgmodel.download(save_dir,1,5)

